import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

def plot_training_curve(log_path, save_path):
    """绘制训练曲线（准确率、损失）"""
    # 从日志文件提取数据（假设日志格式为：epoch, teacher_acc, student_acc, teacher_loss, student_loss）
    epochs = []
    teacher_accs = []
    student_accs = []
    teacher_losses = []
    student_losses = []

    with open(log_path, "r", encoding='utf-8') as f:
        for line in f:
            if "学生训练：损失=" in line and "教师训练：损失=" in line:
                # 解析日志行（需根据实际日志格式调整，此处为示例）
                epoch = int(line.split("第")[1].split("/")[0])
                student_loss = float(line.split("学生训练：损失=")[1].split("，")[0])
                student_acc = float(line.split("学生训练：")[1].split("准确率=")[1].split("，")[0])
                teacher_loss = float(line.split("教师训练：损失=")[1].split("，")[0])
                teacher_acc = float(line.split("教师训练：")[1].split("准确率=")[1].split("，")[0])

                epochs.append(epoch)
                teacher_accs.append(teacher_acc)
                student_accs.append(student_acc)
                teacher_losses.append(teacher_loss)
                student_losses.append(student_loss)

    # 绘制准确率曲线
    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, teacher_accs, label="Teacher Model", marker="o", color="blue")
    plt.plot(epochs, student_accs, label="Student Model", marker="s", color="red")
    plt.xlabel("Epoch")
    plt.ylabel("Training Accuracy")
    # 从log_path提取数据集名称（适配不同数据集，无需写死）
    dataset_name = log_path.split('_')[1]  # 如log_path是"mkd_cifar100_resnet18.log"，提取出"cifar100"
    plt.title(f"MKD Training Accuracy Curve ({dataset_name})")
    plt.legend()
    plt.grid(True)

    # 绘制损失曲线
    plt.subplot(1, 2, 2)
    plt.plot(epochs, teacher_losses, label="Teacher Model", marker="o", color="blue")
    plt.plot(epochs, student_losses, label="Student Model", marker="s", color="red")
    plt.xlabel("Epoch")
    plt.ylabel("Training Loss")
    plt.title(f"MKD Training Loss Curve ({dataset_name})")
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    print(f"训练曲线已保存至：{save_path}")


def plot_attack_advantage(attack_results, dataset_name, save_path):
    """绘制攻击优势率对比图"""
    attack_names = list(attack_results.keys())
    advantages = [attack_results[name][1] for name in attack_names]

    # 论文对比数据（示例，需替换为实际复现的论文指标）
    paper_advantages = {
        "A_corr": 0.121, "A_conf": 0.108, "A_ent": 0.097, "A_mentr": 0.088  # 论文MKD优势率
    }
    paper_advantages = [paper_advantages[name] for name in attack_names]

    # 绘制柱状图
    x = np.arange(len(attack_names))
    width = 0.35

    plt.figure(figsize=(10, 6))
    bars1 = plt.bar(x - width / 2, advantages, width, label="复现MKD", color="skyblue")
    bars2 = plt.bar(x + width / 2, paper_advantages, width, label="论文MKD", color="orange")

    plt.xlabel("成员推理攻击类型")
    plt.ylabel("攻击优势率")
    plt.title(f"{dataset_name}数据集：MKD攻击优势率对比")
    plt.xticks(x, attack_names)
    plt.legend()
    plt.grid(axis="y", linestyle="--")

    # 在柱子上添加数值标签
    for bar in bars1:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2., height + 0.001,
                 f"{height:.4f}", ha="center", va="bottom")
    for bar in bars2:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2., height + 0.001,
                 f"{height:.4f}", ha="center", va="bottom")

    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()
    print(f"攻击优势率对比图已保存至：{save_path}")


# 示例调用（关键修改：所有路径改为绝对路径，避免相对路径问题）
if __name__ == "__main__":
    # -------------------------- 关键修改1：统一使用绝对路径 --------------------------
    # 日志文件绝对路径（确保与train_mkd.py生成的日志路径一致）
    log_abs_path = "E:/xt/雷达/培训/code/results/logs/mkd_cifar100_resnet18.log"
    # 训练曲线保存绝对路径（从log路径提取数据集名称，无需手动写死）
    dataset_name = log_abs_path.split('_')[1]  # 自动提取"cifar100"
    train_curve_save_path = f"E:/xt/雷达/培训/code/results/figures/training_curve_{dataset_name}.png"

    # 绘制训练曲线
    plot_training_curve(
        log_path=log_abs_path,
        save_path=train_curve_save_path
    )

    # -------------------------- 关键修改2：攻击图路径也用绝对路径 --------------------------
    # 绘制攻击优势率对比图（假设attack_results是evaluate_mkd.py的输出，可后续替换为真实结果）
    attack_results = {
        "A_corr": (0.621, 0.121), "A_conf": (0.608, 0.108),
        "A_ent": (0.597, 0.097), "A_mentr": (0.588, 0.088)
    }
    attack_curve_save_path = "E:/xt/雷达/培训/code/results/figures/attack_advantage_cifar100.png"

    plot_attack_advantage(
        attack_results=attack_results,
        dataset_name="CIFAR100",
        save_path=attack_curve_save_path
    )